// <copyright file="EncryptedData.cs" company="SharpSTS">
// Copyright (c) 2007, 2008 All Right Reserved, http://sharpsts.com/
//
// This source is subject to the Microsoft Permissive License.
// Please see the License.txt file for more information.
// All other rights reserved.
//
// THIS CODE AND INFORMATION ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY 
// KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
// PARTICULAR PURPOSE.
//
// </copyright>
// <author>Barry Dorrans</author>
// <email>barryd@idunno.org</email>
// <date>2008-06-10</date>
// <summary>Provides routines for encrypting a token.</summary>

namespace SharpSTS
{
    using System;
    using System.IdentityModel.Selectors;
    using System.IdentityModel.Tokens;
    using System.Runtime.InteropServices;
    using System.Security.Cryptography;
    using System.Xml;

    /// <summary>
    /// A helper class for encrypting security tokens.
    /// </summary>
    internal sealed class EncryptedData
    {
        /// <summary>
        /// Helper class for writing the cipher data elements.
        /// </summary>
        private CipherDataElements cipherData;

        /// <summary>
        /// Helper class for writing the encryption method element.
        /// </summary>
        private EncryptionMethodElement encryptionMethod;

        /// <summary>
        /// Initializes a new instance of the <see cref="EncryptedData"/> class.
        /// </summary>
        public EncryptedData()
        {
            this.cipherData = new CipherDataElements();
            this.encryptionMethod = new EncryptionMethodElement();
        }

        /// <summary>
        /// Sets the encryption algorithm to use.
        /// </summary>
        public string Algorithm
        {
            set
            {
                this.encryptionMethod.Algorithm = value;
            }
        }

        /// <summary>
        /// Gets or sets the key identifier to use.
        /// </summary>
        public SecurityKeyIdentifier KeyIdentifier
        {
            private get;
            set;
        }

        /// <summary>
        /// Gets or sets to serializer to use when serializing the encrypted data to XML.
        /// </summary>
        public SecurityTokenSerializer TokenSerializer
        {
            private get;
            set;
        }

        /// <summary>
        /// Gets or sets the type of XML encryption to use.
        /// </summary>
        public string Type
        {
            private get;
            set;
        }

        /// <summary>
        /// Sets up the encryption process and encrypts the specified buffer.
        /// </summary>
        /// <param name="algorithm">The <see cref="SymmetricAlgorithm"/> to use.</param>
        /// <param name="buffer">The buffer to encrypt.</param>
        /// <param name="offset">The offset into the buffer from which to begin using data.</param>
        /// <param name="length">The number of bytes in the buffer to use as data.</param>
        public void SetUpEncryption(SymmetricAlgorithm algorithm, byte[] buffer, int offset, int length)
        {
            byte[] cipherText;
            byte[] intialisationVector;
            if (null == algorithm)
            {
                throw new ArgumentNullException("algorithm");
            }

            GenerateIVAndEncrypt(algorithm, buffer, offset, length, out intialisationVector, out cipherText);
            this.cipherData.InitialisationVector = intialisationVector;
            this.cipherData.CipherText = cipherText;
        }

        /// <summary>
        /// Writes an XML encryption representation of the cipher data element.
        /// </summary>
        /// <param name="writer">The XML writer to output to.</param>
        public void WriteTo(XmlWriter writer)
        {
            writer.WriteStartElement("enc", "EncryptedData", "http://www.w3.org/2001/04/xmlenc#");
            if (!string.IsNullOrEmpty(this.Type))
            {
                writer.WriteAttributeString("Type", null, this.Type);
            }

            if (!string.IsNullOrEmpty(this.encryptionMethod.Algorithm))
            {
                this.encryptionMethod.WriteTo(writer);
            }

            if (this.KeyIdentifier != null)
            {
                this.TokenSerializer.WriteKeyIdentifier(XmlDictionaryWriter.CreateDictionaryWriter(writer), this.KeyIdentifier);
            }

            this.cipherData.WriteTo(writer);
            writer.WriteEndElement();
        }

        /// <summary>
        /// Generates the intialisation vector and encrypts the specifed plain text.
        /// </summary>
        /// <param name="algorithm">The <see cref="SymmetricAlgorithm"/> to use.</param>
        /// <param name="plainText">The plain text.</param>
        /// <param name="offset">The offset into the byte array from which to begin using data.</param>
        /// <param name="length">The number of bytes in the byte array to use as data.</param>
        /// <param name="intialisationVector">The initialisation vector.</param>
        /// <param name="cipherText">The cipher text.</param>
        private static void GenerateIVAndEncrypt(
            SymmetricAlgorithm algorithm,
            byte[] plainText,
            int offset,
            int length,
            out byte[] intialisationVector,
            out byte[] cipherText)
        {
            RandomNumberGenerator generator = new RNGCryptoServiceProvider();
            int num = algorithm.BlockSize / 8;
            intialisationVector = new byte[num];
            generator.GetBytes(intialisationVector);
            algorithm.Padding = PaddingMode.PKCS7;
            algorithm.Mode = CipherMode.CBC;

            using (ICryptoTransform transform = algorithm.CreateEncryptor(algorithm.Key, intialisationVector))
            {
                cipherText = transform.TransformFinalBlock(plainText, offset, length);
            }
        }

        #region Encryption Structures
        /// <summary>
        /// Represents and writes the encryption information elements.
        /// </summary>
        [StructLayout(LayoutKind.Sequential)]
        public struct CipherDataElements
        {
            /// <summary>
            /// The initialisation vector for the encryption.
            /// </summary>
            internal byte[] InitialisationVector;
            
            /// <summary>
            /// The cipher text to use during encryption.
            /// </summary>
            internal byte[] CipherText;

            /// <summary>
            /// Writes an XML encryption representation of the cipher data element.
            /// </summary>
            /// <param name="writer">The XML writer to output to.</param>
            public void WriteTo(XmlWriter writer)
            {
                writer.WriteStartElement("enc", "CipherData", "http://www.w3.org/2001/04/xmlenc#");
                writer.WriteStartElement("enc", "CipherValue", "http://www.w3.org/2001/04/xmlenc#");
                if (this.InitialisationVector != null)
                {
                    writer.WriteBase64(this.InitialisationVector, 0, this.InitialisationVector.Length);
                }

                writer.WriteBase64(this.CipherText, 0, this.CipherText.Length);
                writer.WriteEndElement();
                writer.WriteEndElement();
            }
        }

        /// <summary>
        /// Represents and writes an element for the XML encryption algorithm used.
        /// </summary>
        [StructLayout(LayoutKind.Sequential)]
        internal struct EncryptionMethodElement
        {
            /// <summary>
            /// The encryption algorithm.
            /// </summary>
            internal string Algorithm;

            /// <summary>
            /// Writes an XML encryption representation of the cipher data element.
            /// </summary>
            /// <param name="writer">The XML writer to output to.</param>
            public void WriteTo(XmlWriter writer)
            {
                writer.WriteStartElement("enc", "EncryptionMethod", "http://www.w3.org/2001/04/xmlenc#");
                writer.WriteAttributeString("Algorithm", null, this.Algorithm);
                writer.WriteEndElement();
            }
        }
        #endregion
    }
}